# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

from absl import app, flags

import pandas
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import pathlib

FLAGS = flags.FLAGS


def plot_histogram(trade_df, trade_change_df, plot_out_root, full_symbol):
  pathlib.Path('%s/hist' % plot_out_root).mkdir(parents=True, exist_ok=True)
  plt.figure(figsize=(10, 10))
  plt.title('histogram of %s trade gap' % full_symbol)
  plt.hist(trade_df['trade_gap'], bins=50)
  plt.yscale('log', basey=2)
  plt.savefig('%s/hist/trade_gap_%s.png' % (plot_out_root, full_symbol))
  plt.close()

  plt.figure(figsize=(10, 10))
  plt.title('histogram of %s trade change gap' % full_symbol)
  plt.hist(trade_change_df['trade_change_gap'], bins=50)
  plt.yscale('log', basey=2)
  plt.savefig('%s/hist/trade_change_gap_%s.png' % (plot_out_root, full_symbol))


def simple_stats_and_print(trade_df, trade_change_df, full_symbol):
  percentiles = [.9, .995]
  print(full_symbol)
  print(trade_df['trade_gap'].describe(percentiles=percentiles))
  print(trade_change_df['trade_change_gap'].describe(percentiles=percentiles))


def cal_trade_change_gap(trade_df):
  trade_change_df = trade_df.groupby(['price', 'qty', 'side'],
                                     as_index=False).first().reset_index(drop=True)
  trade_change_df['trade_change_gap'] = trade_change_df['timestamp'].diff(1)
  trade_change_df['trade_change_gap'] = trade_change_df['trade_gap'].fillna(0)
  return trade_change_df


def cal_trade_gap(trade_df):
  trade_df['trade_gap'] = trade_df['timestamp'].diff(1)
  trade_df['trade_gap'] = trade_df['trade_gap'].fillna(0)
  return trade_df


def read_csv_into_df(csv_path):
  df = pandas.read_csv(csv_path, sep=',', header=0)
  df['datetime'] = pandas.to_datetime(df['timestamp'], unit='ns')
  df = df.set_index('datetime')
  return df


def main(argv):
  csv_path = FLAGS.csv_path
  assert csv_path, '--csv_path must be specified.'
  filename = csv_path.split('/')[-1]
  full_symbol = '.'.join(filename.split('.')[:-1])

  plot_out_root = FLAGS.plot_out_root
  assert plot_out_root, '--plot_out_root must be specified.'

  trade_df = read_csv_into_df(csv_path)
  trade_df = cal_trade_gap(trade_df)
  trade_change_df = cal_trade_change_gap(trade_df)
  simple_stats_and_print(trade_df, trade_change_df, full_symbol)
  plot_histogram(trade_df, trade_change_df, plot_out_root, full_symbol)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('csv_path', None, 'Intput csv file path.')

  flags.DEFINE_string('plot_out_root', None, 'Output plot file root path.')
  app.run(main)
